import os
import numpy as np
from sklearn.datasets import fetch_openml
from sklearn.preprocessing import StandardScaler
import joblib
import logging
from datetime import datetime

logger = logging.getLogger(__name__)


def download_mnist_to_local(dataset='dataset'):
    """
    @brief 下载MNIST数据集到本地指定目录
    @param dataset: 本地保存数据集的目录，默认为 'dataset'

    @details
    该函数使用fetch_openml从OpenML下载MNIST数据集，并将其保存为pkl文件。
    如果目录不存在，则创建该目录。
    下载完成后，记录下载时间和保存路径到日志。

    @exception Exception: 如果下载或保存过程中发生错误，将抛出异常
    @return None
    """
    os.makedirs(dataset, exist_ok=True)

    X, y = fetch_openml('mnist_784', version=1, return_X_y=True, as_frame=False, data_home=dataset)
    y = y.astype(np.int32)
    joblib.dump((X, y), os.path.join(dataset, 'mnist_784.pkl'))

    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    logger.info(f"[{now}] MNIST数据集已保存到 {os.path.join(dataset, 'mnist_784.pkl')}")


def load_mnist_from_internet(train_size=60000, test_size=10000):
    """
    @brief 从网络加载MNIST数据集，并返回标准化后的训练集和测试集
    @param train_size: 训练集大小，默认为60000
    @param test_size: 测试集大小，默认为10000

    @details
    该函数使用fetch_openml从OpenML下载MNIST数据集，并将其划分为训练集和测试集。
    数据集将被标准化处理，以便于后续模型训练和评估。

    @exception Exception: 如果下载或处理过程中发生错误，将抛出异常
    @return X_train, X_test, y_train, y_test: 标准化后的训练集和测试集
    """
    X, y = fetch_openml('mnist_784', version=1, return_X_y=True, as_frame=False)
    y = y.astype(np.int32)

    # 官方划分：前60000为训练集，后10000为测试集
    X_train, X_test = X[:train_size], X[train_size:train_size+test_size]
    y_train, y_test = y[:train_size], y[train_size:train_size+test_size]
    
    # 标准化
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)

    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    logger.info(f"[{now}] MNIST数据集已从网络加载并标准化")
    return X_train, X_test, y_train, y_test


def load_mnist_from_local(data_dir='dataset', train_size=60000, test_size=10000):
    """
    @brief 从本地加载MNIST数据集，并返回标准化后的训练集和测试集
    @param data_dir: 本地数据集目录，默认为 'dataset'
    @param train_size: 训练集大小，默认为60000
    @param test_size: 测试集大小，默认为10000

    @details
    该函数从指定目录加载MNIST数据集的pkl文件，并将其划分为训练集和测试集。
    数据集将被标准化处理，以便于后续模型训练和评估。

    @exception FileNotFoundError: 如果指定路径下的pkl文件不存在，将抛出异常
    @return X_train, X_test, y_train, y_test: 标准化后的训练集和测试集
    """
    path = os.path.join(data_dir, 'mnist_784.pkl')
    if not os.path.exists(path):
        raise FileNotFoundError(f"{path} 不存在，请先运行 download_mnist_to_local()")
    X, y = joblib.load(path)
    y = y.astype(np.int32)

    # 官方划分：前60000为训练集，后10000为测试集
    X_train, X_test = X[:train_size], X[train_size:train_size+test_size]
    y_train, y_test = y[:train_size], y[train_size:train_size+test_size]

    # 标准化
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)

    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    logger.info(f"[{now}] MNIST数据集已从本地加载并标准化")
    return X_train, X_test, y_train, y_test